import torch


#  计算点积的梯度函数
def compute_gradient_dot(x):
    #  定义计算图
    """
    计算 x 与自身的点积：x·x = 0² + 1² + 2² + 3² = 0 + 1 + 4 + 9 = 14
    然后乘以 2：y = 2 * 14 = 28
    """
    y = 2 * torch.dot(x, x)
    print(f"y的值为：{y}")

    # 反向传播计算梯度
    """
    x是一个长度为4的向量，计算x和x的点积，得到了我们赋值给y的标量输出。
    接下来，通过调用反向传播函数来自动计算y关于x每个分量的梯度，并打印这些梯度。
    """
    y.backward()  # 自动计算 y 关于 x 的梯度（导数）,结果会存储在 x.grad 中

    """
    梯度计算原理
    对于 y = 2*(x₀² + x₁² + x₂² + x₃²)，我们计算每个 xᵢ 的偏导数：

    ∂y/∂xᵢ = 2 * 2xᵢ = 4xᵢ

    所以对于 x = [0.0, 1.0, 2.0, 3.0]，梯度应该是：
    [4 * 0.0, 4 * 1.0, 4 * 2.0, 4 * 3.0] = [0.0, 4.0, 8.0, 12.0]
    """
    print(f"x的梯度为：{x.grad}")


#  计算另一个函数的梯度
def compute_gradient_sum(x):
    # 在默认情况下，PyTorch会累积梯度，我们需要清除之前的值
    x.grad.zero_()
    y = x.sum()
    y.backward()
    print(f"计算另一个函数的梯度：{x.grad}")


# 非标量函数的反向传播
def compute_gradient_non_scalar(x):
    # 对非标量调用backward需要传入一个gradient参数，该参数指定微分函数关于self的梯度。
    # 本例只想求偏导数的和，所以传递一个1的梯度是合适的
    x.grad.zero_()
    # 定义非标量 y = [x_0^2, x_1^2, ..., x_n^2]
    y = x * x
    # 方法1：显式求和后反向传播
    """
    y 是一个向量（非标量），直接调用 y.backward() 会报错。
    
    """
    y.sum().backward()
    # 等价于：
    # 方法2：手动传入全1的梯度
    # y.backward(torch.ones(len(x)))


# 分离计算
def compute_gradient_separate_computation(x):
    """
    这段代码演示了 PyTorch 中的 计算图分离（Detach Computation） 技术，
    主要用于控制梯度的传播范围。
    """
    # x.grad.zero_()
    """
    对 x 逐元素平方，得到 y。
    例如，若 x = [0, 1, 2, 3]，则 y = [0, 1, 4, 9]。
    关键点：y 是计算图的一部分，其梯度会回溯到 x。
    """
    y = x * x
    # y.detach() 返回一个与 y 值相同的新张量 u，但 脱离计算图。
    """
    u 的特性：
        * 值与 y 相同（如 [0, 1, 4, 9]）。
        * 不携带梯度信息（requires_grad=False）。
        * 对 u 的操作不会影响 y 或 x 的梯度计算。
    打印 u 的类型会显示 <class 'torch.Tensor'>，但无梯度追踪。
    """
    u = y.detach()
    print(f"u={u},类型是：{type(u)}")
    """
    将分离后的 u 与原始 x 逐元素相乘。
    例如，z = [0 * 0, 1 * 1, 4 * 2, 9 * 3] = [0, 1, 8, 27]。
    关键点：
        * 由于 u 被分离，乘法操作 不会通过y回溯到x的梯度。
        * 梯度仅通过第二个 x 传播（即 z 对 x 的梯度仅来自 x 的直接参与）。
    """
    z = u * x
    z.sum().backward()
    print(f"分离计算的梯度为：{x.grad}")


if __name__ == "__main__":
    # 创建张量并启用梯度追踪
    x = torch.arange(4.0)
    x.requires_grad_(True)  # 等价于x=torch.arange(4.0,requires_grad=True)
    # print(x.grad)  # x.grad为None

    # 计算点积
    # compute_gradient_dot(x)

    #  计算另一个函数的梯度
    # compute_gradient_sum(x)

    # 非标量函数的反向传播
    # compute_gradient_non_scalar(x)

    # 分离计算
    compute_gradient_separate_computation(x)
